import asyncio
import json
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
import uuid
from typing import Dict, AsyncGenerator
import time

# 全局变量用于存储客户端连接
clients: Dict[str, asyncio.Queue] = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时初始化
    yield
    # 关闭时清理所有客户端连接
    clients.clear()


app = FastAPI(lifespan=lifespan)

# 添加CORS中间件以允许前端访问
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应限制为具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


async def event_generator(client_id: str) -> AsyncGenerator[str, None]:
    """
    为指定客户端生成事件流的异步生成器
    """
    try:
        # 先发送一条连接成功的消息
        yield f"data: {json.dumps({'type': 'connected', 'client_id': client_id, 'message': 'Connection established'})}\n\n"

        while True:
            # 从客户端的队列中获取数据
            if client_id not in clients:
                break

            try:
                # 设置超时以避免永久阻塞
                data = await asyncio.wait_for(clients[client_id].get(), timeout=30.0)
                if data is None:  # 终止信号
                    break

                # 格式化数据为SSE格式
                yield f"data: {json.dumps(data)}\n\n"
            except asyncio.TimeoutError:
                # 发送心跳保持连接
                yield f"data: {json.dumps({'type': 'heartbeat', 'timestamp': time.time()})}\n\n"
                continue
    except asyncio.CancelledError:
        print(f"Client {client_id} disconnected")
    finally:
        # 清理客户端连接
        if client_id in clients:
            del clients[client_id]


@app.get("/stream")
async def stream_data(request: Request):
    """
    处理前端GET请求并返回流式响应
    EventSource只能使用GET请求，所以改为GET方法
    """
    # 生成唯一客户端ID
    client_id = str(uuid.uuid4())

    # 为客户端创建消息队列
    clients[client_id] = asyncio.Queue()

    # 返回流式响应
    response = StreamingResponse(
        event_generator(client_id),
        media_type="text/event-stream",
    )

    # 设置SSE所需的响应头
    response.headers["Cache-Control"] = "no-cache"
    response.headers["Connection"] = "keep-alive"
    response.headers["Access-Control-Allow-Origin"] = "*"

    return response


@app.post("/push/{client_id}")
async def push_data(client_id: str, request: Request):
    """
    接收其他Python脚本发送的流式数据并转发给指定客户端
    """
    if client_id not in clients:
        return {"status": "error", "message": f"Client not found: {client_id}"}

    try:
        # 异步读取请求体
        body = await request.body()
        data = json.loads(body.decode("utf-8"))

        # 将数据放入客户端的队列
        await clients[client_id].put(data)

        return {"status": "success", "message": "Data pushed to client"}
    except Exception as e:
        return {"status": "error", "message": str(e)}


@app.get("/clients")
async def list_clients():
    """
    列出所有活跃客户端（用于调试）
    """
    return {"clients": list(clients.keys())}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)